from air_web.data_platform import sql_engine


def get_cons_num():
    sql = """select is_cal, count(1) as cons_num
             from c_cons
             where c_id < 100000000
             group by is_cal
          """
    df = sql_engine.query(sql)
    num_dict = df.set_index("is_cal")["cons_num"].to_dict()
    cal_cons_num = num_dict[1]
    no_cal_cons_num = num_dict[0] if 0 in num_dict else 0
    all_cons_num = cal_cons_num + no_cal_cons_num
    return all_cons_num, cal_cons_num, no_cal_cons_num


def get_avg_p_total_from_mysql():
    sql = """select org_no, org_name, data_time, avg_p_total, avg_p_kt
             from orgno_typeid_day_stat
             where org_no in (select org_no from real_org_no where org_level<2)
               and type_id=0
               and data_time >= (
               SELECT MAX(data_time) - INTERVAL 6 DAY FROM orgno_typeid_day_stat
               where org_no in (select org_no from real_org_no where org_level=0)
               and type_id=0
               )
             order by org_no, data_time
          """
    df = sql_engine.query(sql)
    df.loc[df["org_name"].isna(), "org_name"] = "全省"
    return df


if __name__ == "__main__":
    get_avg_p_total_from_mysql()
    get_cons_num()
